import cv2
import os
import xml
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


def decide_order(points):
    points = sorted(points, key=lambda x:x[0])
    first_and_last_point = sorted(points[:2], key=lambda x:x[1])
    second_and_third_point = sorted(points[-2:], key=lambda x:x[1])
    points = np.concatenate([first_and_last_point[0:1], second_and_third_point, first_and_last_point[1:]], axis=0)
    return points

def parser_pascal_voc_xml(xml_path):
    import xml.etree.ElementTree as ET
    import logging
    image_path = os.path.join(xml_path[:-4] + ".jpg")
    if cv2.imread(image_path) is None:
        return None, None
    image = cv2.imread(image_path)[:, :, ::-1]
    image = cv2.GaussianBlur(image, ksize=(15, 15), sigmaX=1, sigmaY=1)
    oneimg = {}
    oneimg['bndbox'] = []
    try:
        dom = ET.parse(xml_path)
    except Exception as e:
        return None, None
    root = dom.getroot()
    filename = root.findall('filename')[0].text
    # oneimg['path'] = os.path.join(img_root, filename)
    oneimg['filename'] = filename
    image_cropped_list = [None, None]
    for objects in root.findall('object'):
        name = objects.find('name').text
        points = list(objects.find('polygon'))
        if len(points) != 4:
            break
        point_list = [x.text.strip().split(',') for x in points]
        points = np.array(point_list).astype(np.int32)
        points = decide_order(points)
        points_dst = np.array([[0, 0], [24, 0], [24, 128], [0, 128]])
        M = cv2.estimateAffine2D(points, points_dst, method=cv2.LMEDS)[0]
        image_cropped = cv2.warpAffine(image, M, dsize=(24, 128))
        image_cropped = image_cropped[:, 8:17]
        image_cropped = np.array(Image.fromarray(image_cropped).convert("CMYK"))[:, :, 2]
        cv2.imshow("ss" + name, image_cropped)
        cv2.waitKey(1)
        image_cropped2 = (( image_cropped > 128).astype('f').sum(axis=1)) > 4
        image_cropped = image_cropped2[:, np.newaxis]
        image_cropped_list[["1111", "2222"].index(name)] = image_cropped

    if image_cropped_list[0] is None or image_cropped_list[1] is None:
        return None, None
    return image, image_cropped_list


def parse_labelme(image, json_path):
    import json
    obj = json.load(open(json_path, "rb"))
    shapes = obj["shapes"]
    points = {}
    for shape in shapes:
        if shape["shape_type"]=="point":
            points[shape["label"]] = [shape["points"][0][0], shape["points"][0][1]]
    points_src = np.array([points["a"], points["b"], points["c"], points["d"]])
    points_dst = np.array([[0, 0], [24, 0], [24, 128], [0, 128]])

    M = cv2.findHomography(points_src, points_dst, method=cv2.LMEDS)
    # image_cropped = cv2.warpAffine(image, M[0], dsize=(24, 128))
    image_cropped = cv2.warpPerspective(image, M[0], dsize=(24, 128))
    image_cropped = image_cropped[:, 8:17]
    return image_cropped


def recog_in_image_cropped(image_cropped_o):
    image_cropped_blur = cv2.GaussianBlur(image_cropped_o, ksize=(3, 3), sigmaX=0, sigmaY=1)
    image_cropped_blur = cv2.GaussianBlur(image_cropped_blur, ksize=(3, 3), sigmaX=0, sigmaY=1)
    image_cropped_blur = cv2.GaussianBlur(image_cropped_blur, ksize=(3, 3), sigmaX=0, sigmaY=1)

    image_cropped = image_cropped_blur.astype("f")

    image_cropped = image_cropped[:, :, 0] - image_cropped[:, :, 1]
    binary0 = ((image_cropped > 15).astype('f').sum(axis=1)) > 4

    fig, axes = plt.subplots(1, 2)
    axes[0].imshow(binary0[:, np.newaxis])
    axes[1].imshow(image_cropped_o)

    first_level = 0
    for i in range(0, binary0.shape[0] - 4, 1):
        h = binary0.shape[0]
        if sum(binary0[h - i - 4:h - i]) == 0:
            first_level = i
            break
    print("first level", first_level / h)

    plt.show()


if __name__ == '__main__':

    for root,_, names in os.walk("/home/kohill/Desktop/仪表识别"):
        for name in names:
            path = os.path.join(root, name)
            if "json" in path:
                assert os.path.exists(path)
                image_path = str(path).replace(".json", ".jpg")
                image = cv2.imread(image_path)[:, :, ::-1]
                image_cropped = parse_labelme(image, path)
                recog_in_image_cropped(image_cropped)

            elif ".xml" in path:
                image, image_cropped_list = parser_pascal_voc_xml(path)

                if image_cropped_list is None:
                    continue
                axes[0].imshow(image_cropped_list[0])
                axes[1].imshow(image_cropped_list[1])
                binary0 = image_cropped_list[0].squeeze()
                binary1 = image_cropped_list[1].squeeze()
                # find the level of the first scale
                first_level = 0
                for i in range(0, binary0.shape[0]-4, 1):
                    h = binary0.shape[0]
                    if sum(binary0[h-i-4:h-i]) == 0:
                        first_level = i
                        break
                print("first level", first_level)
                # find the level of the second scale

                second_level = 0
                for i in range(0, binary1.shape[0]-16, 1):
                    h = binary1.shape[0]
                    if sum(binary1[i:i+16]) == 16:
                        second_level = i
                        break
                print("second level", second_level)

                cv2.imshow("", cv2.resize(image, (0,0), fx=.5, fy=0.5)[:, :, ::-1])
                cv2.waitKey(0)
